# !pip install scikit-image
# !pip install --upgrade gradio
# !pip install requests opencv-python

import cv2
import numpy as np
from skimage.morphology import skeletonize
# import gradio as gr
import requests

def process_image(image):
    result = {}

    # Convert image to grayscale
    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # 二值化处理
    _, binary = cv2.threshold(gray_image, 240, 255, cv2.THRESH_BINARY)

    # 检测轮廓
    contours, _ = cv2.findContours(binary, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    # 创建一个空白图像用于绘制检测结果
    output = cv2.cvtColor(gray_image, cv2.COLOR_GRAY2BGR)

    # 绘制轮廓
    # 1则 线条变细，2线条加粗
    cv2.drawContours(output, contours, -1, (0, 255, 0), 1)

    # 将结果图像保存到内存中
    _, buffer = cv2.imencode('.jpg', output)
    detected_contours = cv2.imdecode(np.frombuffer(buffer, np.uint8), cv2.IMREAD_COLOR)

    # 转换为灰度图像
    gray_wall_only2 = cv2.cvtColor(detected_contours, cv2.COLOR_BGR2GRAY)

    # 二值化处理
    _, binary_image2 = cv2.threshold(gray_wall_only2, 127, 255, cv2.THRESH_BINARY)

    # 骨架化
    skeleton2 = skeletonize(binary_image2 // 255)
    skeleton2 = (skeleton2 * 255).astype(np.uint8)
    skeleton_image2 = np.zeros_like(detected_contours)
    skeleton_image2[skeleton2 == 255] = [255, 255, 255]

    # 检测骨架图像的轮廓
    skeleton_contours2, _ = cv2.findContours(skeleton2, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
    linedoor = []
    for i, contour in enumerate(skeleton_contours2):
        for j in range(len(contour)):
            x1, y1 = contour[j][0]
            x2, y2 = contour[(j + 1) % len(contour)][0]
            linedoor.append((1, x1, y1, x2, y2))

    # Detect white and black regions
    lower_white = np.array([200, 200, 200])
    upper_white = np.array([255, 255, 255])
    white_mask = cv2.inRange(image, lower_white, upper_white)

    lower_black = np.array([0, 0, 0])
    upper_black = np.array([50, 50, 50])
    black_mask = cv2.inRange(image, lower_black, upper_black)

    # Find contours for windows and walls
    window_contours, _ = cv2.findContours(white_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    wall_contours, _ = cv2.findContours(black_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # Detect edges for door using Canny edge detection
    threshold1 = 50
    threshold2 = 150
    edges = cv2.Canny(image, threshold1, threshold2)
    kernel = np.ones((5, 5), np.uint8)
    edges = cv2.dilate(edges, kernel, iterations=1)
    edges = cv2.erode(edges, kernel, iterations=1)
    contours, _ = cv2.findContours(edges, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)

    door_contours = []
    for contour in contours:
        x, y, w, h = cv2.boundingRect(contour)
        perimeter = cv2.arcLength(contour, True)
        approx = cv2.approxPolyDP(contour, 0.04 * perimeter, True)
        if len(approx) > 3 and 0.5 < w / float(h) < 2 and 600 < cv2.contourArea(contour) < 1500:
            door_contours.append(contour)

    # Morphological operations and re-detect contours
    white_mask = cv2.morphologyEx(white_mask, cv2.MORPH_OPEN, kernel)
    white_mask = cv2.morphologyEx(white_mask, cv2.MORPH_CLOSE, kernel)
    black_mask = cv2.morphologyEx(black_mask, cv2.MORPH_OPEN, kernel)
    black_mask = cv2.morphologyEx(black_mask, cv2.MORPH_CLOSE, kernel)
    window_contours, _ = cv2.findContours(white_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    wall_contours, _ = cv2.findContours(black_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # Create door-only image
    door_only_image = np.zeros_like(image)
    cv2.drawContours(door_only_image, door_contours, -1, (0, 255, 0), 1)

    # Skeletonize door area
    gray_wall_only2 = cv2.cvtColor(door_only_image, cv2.COLOR_BGR2GRAY)
    _, binary_image2 = cv2.threshold(gray_wall_only2, 127, 255, cv2.THRESH_BINARY)
    skeleton2 = skeletonize(binary_image2 // 255)
    skeleton2 = (skeleton2 * 255).astype(np.uint8)
    skeleton_image2 = np.zeros_like(image)
    skeleton_image2[skeleton2 == 255] = [255, 255, 255]
    skeleton_contours2, _ = cv2.findContours(skeleton2, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
    linedoor = []
    for i, contour in enumerate(skeleton_contours2):
        for j in range(len(contour)):
            x1, y1 = contour[j][0]
            x2, y2 = contour[(j + 1) % len(contour)][0]
            linedoor.append((1, x1, y1, x2, y2))

    # Skeletonize wall area
    # 窗户多线条
    # gray_wall_only = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    # 窗户单线条
    gray_wall_only = cv2.cvtColor(detected_contours, cv2.COLOR_BGR2GRAY)
    _, binary_image = cv2.threshold(gray_wall_only, 127, 255, cv2.THRESH_BINARY)
    skeleton = skeletonize(binary_image // 255)
    skeleton = (skeleton * 255).astype(np.uint8)
    skeleton_image = np.zeros_like(image)
    skeleton_image[skeleton == 255] = [255, 255, 255]
    skeleton_contours, _ = cv2.findContours(skeleton, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
    linewall = []
    for i, contour in enumerate(skeleton_contours):
        for j in range(len(contour)):
            x1, y1 = contour[j][0]
            x2, y2 = contour[(j + 1) % len(contour)][0]
            linewall.append((0, x1, y1, x2, y2))

    result['skeleton_image'] = skeleton_image
    result['skeleton_image2'] = skeleton_image2
    result['wall_contours'] = linewall
    result['door_contours'] = linedoor

    return result

def download_image(url):
    response = requests.get(url)
    image_array = np.asarray(bytearray(response.content), dtype=np.uint8)
    image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
    return image

def gradio_interface(image):
    processed_result = process_image(image)
    # return processed_result['skeleton_image'], processed_result['skeleton_image2'], str(processed_result['wall_contours']), str(processed_result['door_contours'])
    return processed_result['skeleton_image'], processed_result['skeleton_image2'], processed_result['wall_contours'], processed_result['door_contours']

# 在notebook中显示图像
def display_image(image, title='Image'):
    plt.figure(figsize=(6, 6))
    if len(image.shape) == 2:  # 灰度图像
        plt.imshow(image, cmap='gray')
    else:  # 彩色图像
        plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    plt.title(title)
    plt.axis('off')
    plt.show()

# with gr.Blocks() as demo:
#     with gr.Tab("贝壳户型图识别"):
#         with gr.Row():
#             imageSource = gr.Image(label="贝壳户型图", sources=["upload"])
#             imagewallTarget = gr.Image(label="贝壳线墙线框图")
#             imagedoorTarget = gr.Image(label="贝壳线门体框图")
#         with gr.Row():
#             wallTarget = gr.Textbox(label="墙体json", container=False,  lines = 200)
#             doorTarget = gr.Textbox(label="门体json", container=False,  lines = 200)
        
#         imageSource.change(gradio_interface, inputs=imageSource, outputs=[imagewallTarget, imagedoorTarget, wallTarget, doorTarget])

# if __name__ == "__main__":
#     demo.launch()
#     # demo.launch(share=True, debug=True, server_name="mangoai.dev.com", server_port=8820)

import matplotlib.pyplot as plt
# 示例URL
# 原图https://ke-image.ljcdn.com/hdic-frame/standard_014c0ced-3659-405f-b096-663dc56e067f.png!m_fill,w_1000,h_750?from=ke.com
# image_url = 'https://vrlab-public.ljcdn.com/release/vrcustomer/outline_0_014c0ced-3659-405f-b096-663dc56e067f.png'

# image_url = 'https://vrlab-public.ljcdn.com/release/vrcustomer/outline_0_477d6cbd-9d4f-42cf-b819-4082813fc63a.png'
# image_url = 'https://vrlab-public.ljcdn.com/release/vrcustomer/outline_0_8a2bcc56-4873-43d1-b48c-9791cab1ce35.png'
# image_url = 'https://vrlab-public.ljcdn.com/release/vrcustomer/outline_0_ddfb1e2f-564b-46de-9e05-73c72c408d12.png'
# image_url = 'https://vrlab-public.ljcdn.com/release/vrcustomer/outline_0_761bf2a5-6823-4d5d-b842-ab0e03589057.png'
image_url = 'https://vrlab-public.ljcdn.com/release/vrcustomer/outline_0_01bc86ca-e8d9-44e8-b92f-d5ba27d6cb72.png'
# image_url = 'https://vrlab-public.ljcdn.com/release/vrcustomer/outline_0_96bd87f6-d5fd-4fa0-b300-d96ed7cb0c89.png'

# 下载图片
image = download_image(image_url)

# 调用process_image函数
result = gradio_interface(image)
# print(result)
# display_image(result[0], 'Original Image')
# display_image(result[1], 'Original Image')
# print(result[2],result[3])

xy_bboxes = result[2]
door_data = result[3]

# 转到 FloorplanImageDataConverter.py __main__
from FloorplanImageDataConverter import FloorplanImageDataConverter

fl_img_dc = FloorplanImageDataConverter()
# fl_img_dc.showImg(door_data,[])
# fl_img_dc.showImg(xy_bboxes,[]) 

#  墙是0，门是1，窗户是3
all_data = xy_bboxes + door_data

# todo: 获取比例信息     
radio_wall = 24.6767
# radio_wall = 1

# fl_img_dc.showImg(all_data,[],False)
# fl_img_dc.showImg(xy_bboxes,[]) 
x_data,y_data = fl_img_dc.run(all_data, radio_wall)
fl_img_dc.showImg(x_data,y_data,False)

# 转换
data_c = (x_data + y_data)
# print(data_c)
from KeFloorplanConverter import KeFloorplanConverter
import json


keFloorplanConverter = KeFloorplanConverter()
mangoDict = keFloorplanConverter.run(data_c)
print(json.dumps(mangoDict, ensure_ascii=False))